import time
import json
import cv2
import math
from tool import has_digit_input
from cap import VideoCapture
from arm_module import Arm
from TPU import TPUFactory
from get_position import RectangleDetector

class FruitDetector:
    def __init__(self, config_file='model/FD.conf', display=True):
        self.display = display
        self.config_file = config_file
        self._load_config()
        self._initialize_variables()
        self.running = False
        self.obstacle_mode = False
        self.obstacle_point = None
        
    def _load_config(self):
        """加载配置文件"""
        with open(self.config_file, 'r', encoding='utf-8') as f:
            self.config = json.load(f)
            
        self.model_id = self.config['model_id']
        self.classes = self.config['classes']
        self.confThres = self.config['conf']['confThres']
        self.iouThres = self.config['conf']['iouThres']
        self.imgsz = self.config['conf']['imgsz']
        self.model_version = self.config['model_version']
        self.modelType = self.config['modelType']
        
    def _initialize_variables(self):
        """初始化变量"""
        self.fruit_classes = {
            "defect": 0,
            "freshApple": 1,
            "freshBanana": 2,
            "freshOrange": 3,
            "mold": 4,
            "rottenApple": 5,
            "rottenBanana": 6,
            "rottenOrange": 7
        }

        # 水果计数器
        self.cnt = [0] * len(self.classes)
        self.TIMES = 1  # 识别次数阈值
        
        # 初始化设备
        self.arm = None
        self.tpu = None
        self.cap = None
        
    def initialize_devices(self):
        """初始化所有设备"""
        try:
            # 初始化TPU
            self.tpu = TPUFactory(algorithm="detection")
            if not self.tpu.link():
                raise Exception("初始化TPU异常")

            # 初始化模型
            model_params = {
                "model_name": "FD.lfdata",
                "model_conf_name": "FD.conf",
                "num_classes": len(self.classes),
                "model_size": (3, 640, 640),
                "ai_app_name": "det_app",
                "model_version": self.model_version,
                "modelType": self.modelType
            }
            print("准备初始化模型...")
            self.tpu.model_init(model_params)
            
            return True
            
        except Exception as e:
            print(f"初始化设备时出错: {e}")
            self.release_resources()
            return False
            
    def open_camera(self):
        """打开摄像头"""
        try:
            if self.cap is None:
                print("准备初始化摄像头...")
                self.cap = VideoCapture()
                return True
            return False
        except Exception as e:
            print(f"打开摄像头时出错: {e}")
            return False
            
    def close_camera(self):
        """关闭摄像头"""
        if hasattr(self, 'cap') and self.cap is not None:
            self.cap.release()
            self.cap = None
            cv2.destroyAllWindows()
            return True
        return False
            
    def start_detection(self):
        """开始检测"""
        # if not self.initialize_devices():
            #exit(1)
        if not all([self.tpu, self.cap]):
            print("请先初始化设备和打开摄像头")
            return False
            
        self.running = True
        return True
        
    def stop_detection(self):
        """停止检测"""
        self.running = False
        # if hasattr(self, 'tpu') and self.tpu is not None:
            # self.tpu.release()
        return True
            
    def get_frame(self):
        """获取当前摄像头帧"""
        if hasattr(self, 'cap') and self.cap is not None:
            return self.cap.read()
        return None
    
    def detect_fruits(self, img):
        """检测水果并返回检测结果"""
        if img is None:
            return None, None
            
        boxes = self.tpu.predict(img, self.imgsz, self.confThres, self.iouThres)
        result_img = img.copy()
        
        if boxes:
            # 更新计数器
            for o in boxes:
                label = int(o[5])
                self.cnt[label] += 1
                
            # 绘制检测结果
            for o in boxes:
                x = int(o[0])
                y = int(o[1])
                w = int(o[2])
                h = int(o[3])
                score = o[4]
                label = int(o[5])
                text = f"{self.classes[label]}:{score:.2f}"
                
                # 计算文本位置（上边框中央）
                text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_PLAIN, 1.5, 2)[0]
                text_x = x - text_size[0] // 2
                text_y = y - h // 2 + 30  # 上边框上方10像素
                cv2.putText(result_img, text, (text_x, text_y), cv2.FONT_HERSHEY_PLAIN, 1.5, (255, 0, 0), thickness=2)
                cv2.rectangle(result_img, (x - w//2, y - h//2), (x + w//2, y + h//2), (255, 0, 0), 3)
                # 在物体中心位置绘制圆点(蓝色)
                cv2.circle(result_img, (x, y), 5, (255, 0, 0), -1)
        
        # 绘制障碍物点
        if self.obstacle_point:
            cv2.circle(result_img, (self.obstacle_point[0], self.obstacle_point[1]), 10, (0, 0, 255), -1)
        
        if self.display:
            cv2.imshow("video", result_img)
            cv2.waitKey(1)
            
        return boxes, result_img
        
    def should_grab(self, boxes):
        """判断是否需要抓取"""
        if not boxes:
            # 重置计数器
            for i in range(len(self.cnt)):
                self.cnt[i] = 0
            return None
            
        for o in boxes:
            label = int(o[5])
            if self.cnt[label] >= self.TIMES:
                self.cnt[label] = 0  # 重置计数器
                return o  # 返回需要抓取的物体
        
        return None
        
    def run_detection_loop(self):
        """运行检测循环"""
        try:
            while self.running:
                img = self.get_frame()
                boxes, result_img = self.detect_fruits(img)
                
                if has_digit_input() == 'q':
                    break

        finally:
            self.stop_detection()
            print("检测循环结束")
            
    def release_resources(self):
        """释放所有资源"""
        self.stop_detection()
        self.close_camera()
        # cv2.destroyAllWindows()

    def set_obstacle_mode(self, mode):
        """设置障碍物模式"""
        self.obstacle_mode = mode
        if not mode:
            self.obstacle_point = None

    def set_obstacle_point(self, point):
        """设置障碍物点"""
        self.obstacle_point = point


class FruitGrabber:
    def __init__(self):
        self.arm = None
        self.targetPoints = [
            [150, -130],  # 放置新鲜苹果
            [150, -240],  # 放置新鲜香蕉
            [-150, -240],  # 放置新鲜橙子
            [-150, -130]  # 放置腐烂水果
        ]
        self.HEIGHT = 130
        self.obstacle_point = None
        
    def initialize_arm(self):
        """初始化机械臂"""
        try:
            self.arm = Arm()
            self.arm.reset_position()
            return True
        except Exception as e:
            print(f"初始化机械臂时出错: {e}")
            return False
            
    def grab_fruit(self, detection_result, label, obstacle_mode=False):
        """抓取水果"""
        #if not self.arm:
            #print("机械臂未初始化")
            #return False
        if not self.arm:
            print("正在初始化机械臂")
            self.arm = Arm()
            self.arm.reset_position()
            print("已初始化机械臂")
            #return False
            
        x = int(detection_result[0])
        y = int(detection_result[1])
        
        if label in [1, 2, 3]:  # 新鲜水果
            target_point = self.targetPoints[label - 1]
        else:  # 腐烂水果
            target_point = self.targetPoints[3]
            
        if obstacle_mode and self.obstacle_point:
            self.move_obstacle(x, y, target_point[0], target_point[1])
        else:
            self._execute_grab_sequence(x, y, target_point[0], target_point[1])
        return True
        
    def _execute_grab_sequence(self, pixelX, pixelY, targetX, targetY):
        """执行抓取序列"""
        up_height = 180
        self.arm.suck_up()
        
        # 下降到吸取高度
        print(f"移动到卡片位置: x={pixelX}, y={pixelY}, z={self.HEIGHT}")
        times, coordinate = self.arm.move_pixel(pixelX, pixelY, self.HEIGHT)
        time.sleep(1)
        
        # 抬起
        print(f"抬起到位置: x={coordinate.x}, y={coordinate.x}, z={up_height}")
        self.arm.move_point(coordinate.x, coordinate.y, up_height, 1500)
        
        # 移动到放置位置上方
        print(f"移动到位置: x={targetX}, y={targetY}, z={up_height}")
        self.arm.move_point(targetX, targetY, up_height, 1500)
        
        # 下降到放置高度
        print(f"下降到位置: x={targetX}, y={targetY}, z={135}")
        self.arm.move_point(targetX, targetY, 135, 1500)
        self.arm.suck_release()
        time.sleep(1)
        
        # 移动到放置位置上方
        print(f"抬起到位置: x={targetX}, y={targetY}, z={up_height}")
        self.arm.move_point(targetX, targetY, up_height, 1500)
        
        # 移动到中间位置
        middle_x = (targetX + coordinate.x)/2
        middle_y = (targetY + coordinate.y)/2
        print(f"抬起到位置: x={middle_x}, y={middle_y}, z={up_height}")
        self.arm.move_point(middle_x, middle_y, up_height, 800)
        
        self.arm.move_point(0, -115, 200, 800)
        self.arm.move_point(0, -115, 240, 800)
        time.sleep(0.8)
        
    def move_obstacle(self, pixelX, pixelY, x, y):
        """机械臂规避障碍物进行抓取函数"""
        up_height = 180
        obstacle_height = 200  # 障碍物高度
        self.arm.suck_up()
        
        # 获取障碍物坐标
        obstacle_x, obstacle_y = self.obstacle_point
        obstacle_x = int(obstacle_x)  # 转换为 Python int
        obstacle_y = int(obstacle_y)  # 转换为 Python int
        
        # 打印障碍物坐标的数据类型
        print(f"obstacle_x 的数据类型: {type(obstacle_x)}")
        print(f"obstacle_y 的数据类型: {type(obstacle_y)}")
       
        # 1. 移动到卡片位置
        print(f"移动到卡片位置: x={pixelX}, y={pixelY}, z={self.HEIGHT}")
        times, coordinate = self.arm.move_pixel(pixelX, pixelY, self.HEIGHT)
        
        # 2. 抬起
        print(f"抬起到位置: x={coordinate.x}, y={coordinate.x}, z={up_height}")
        self.arm.move_point(coordinate.x, coordinate.y, up_height, 1500)

        # 3. 移动到障碍物的上方位置
        print(f"移动到障碍物的上方位置: x={obstacle_x}, y={obstacle_y}, z={obstacle_height}")
        usingtime, obstacle_coordinate = self.arm.move_pixel(obstacle_x, obstacle_y, obstacle_height)
        
        # 4. 绕过障碍物
        bypass_x = (obstacle_coordinate.x + x)/2  #  绕过障碍物时的X坐标
        bypass_y = (obstacle_coordinate.y + y)/2  # 绕过障碍物时的Y坐标
        print(f"绕过障碍物: x={bypass_x}, y={bypass_y}, z={obstacle_height}")
        self.arm.move_point(bypass_x, bypass_y, obstacle_height, 1500)
        
        # 5. 移动到放置位置上方
        print(f"移动到放置位置上方: x={x}, y={y}, z={up_height}")
        self.arm.move_point(x, y, up_height, 1500)
        
        # 6. 下降到放置高度
        print(f"下降到放置高度: x={x}, y={y}, z={135}")
        self.arm.move_point(x, y, 135, 1500)
        self.arm.suck_release()
        time.sleep(1)
        
        # 7. 移动到放置位置上方
        print(f"抬起到安全高度: x={x}, y={y}, z={up_height}")
        self.arm.move_point(x, y, up_height, 1500)

        # 8. 绕过障碍物
        print(f"绕过障碍物: x={bypass_x}, y={bypass_y}, z={obstacle_height}")
        self.arm.move_point(bypass_x, bypass_y, obstacle_height, 1500)

        # 9. 移动到障碍物的上方位置
        print(f"移动到障碍物的上方位置: x={obstacle_coordinate.x}, y={obstacle_coordinate.y}, z={obstacle_height}")
        self.arm.move_point(obstacle_coordinate.x,obstacle_coordinate.y, obstacle_height,1500)

        # 10. 复位
        self.arm.reset_position()
        time.sleep(0.8)
        
    def reset_position(self):
        """复位机械臂"""
        if self.arm:
            self.arm.reset_position()
            
    def release_resources(self):
        """释放资源"""
        if self.arm:
            self.arm.reset_position()
            
    def set_obstacle_point(self, point):
        """设置障碍物点"""
        self.obstacle_point = point



class FruitSortingSystem:
    def __init__(self):
        self.detector = FruitDetector(display=False)
        self.grabber = FruitGrabber()
        self.obstacle_mode = False
        self.rectangle_detector = RectangleDetector()  # 新增矩形检测器
        
    def start(self):
        """启动系统"""
        # 初始化摄像头
        # if not self.detector.open_camera():
            # return False
            
        # 初始化机械臂
        if not self.grabber.initialize_arm():
            return False
            
        # 初始化检测器(TPU和模型)
        if not self.detector.initialize_devices():
            return False
        
        # 开始检测
        # if not self.detector.start_detection():
            # return False
            
        return True
        
    def set_obstacle_from_rectangle(self,img):
        """通过矩形检测设置障碍物点"""
        print("正在检测矩形以设置障碍物点...")
        point = self.rectangle_detector.detect_rectangles(img)
        if point:
            self.detector.set_obstacle_point(point)
            self.grabber.set_obstacle_point(point)
            print(f"从矩形检测设置障碍物点: {point}")
            return True
        print("未检测到矩形，无法设置障碍物点")
        return False
        
    def run(self):
        """运行主循环"""
        self.detector.start_detection()
        try:
            while self.detector.running:                
                # 开始检测循环
                while True:
                    print("请放置...")
                    time.sleep(5)
                    
                    img = self.detector.get_frame()
                    if img is None:
                        continue
                    
                    # 询问是否设置障碍物
                    user_input = input("是否设置障碍物? (y/n/q): ").lower()
                    
                    if user_input == 'y':
                        self.obstacle_mode = True
                        self.detector.set_obstacle_mode(True)
                        # 通过矩形检测设置障碍物点
                        if not self.set_obstacle_from_rectangle(img):
                            continue
                    elif user_input == 'n':
                        self.obstacle_mode = False
                        self.detector.set_obstacle_mode(False)
                        print("已设置为正常模式")
                        
                    boxes, _ = self.detector.detect_fruits(img)
                    
                    # 检查是否需要抓取
                    grab_target = self.detector.should_grab(boxes)
                    if grab_target:
                        label = int(grab_target[5])
                        self.grabber.grab_fruit(grab_target, label, self.obstacle_mode)
                    
        except KeyboardInterrupt:
            print("用户中断")
        finally:
            self.stop()
            
    def stop(self):
        """停止系统"""
        self.detector.stop_detection()
        self.grabber.reset_position()
        self.detector.release_resources()
        self.grabber.release_resources()
        




if __name__ == "__main__":
    system = FruitSortingSystem()
    if system.start():
        system.run()
